package cc.xpcas.raft.context;

import java.util.concurrent.*;
import java.util.function.*;

import lombok.Setter;

/**
 * @author xp
 */
public class ScheduledFuture<V> implements java.util.concurrent.RunnableFuture<V> {

    private final Object lock = new Object();

    private volatile int state = 0; // 0 = not started, 1 = running, 2 = done, 3 = cancelled

    private final Supplier<V> supplier;

    private V result = null;

    private Throwable throwable = null;

    private Thread runningOn;

    @Setter
    private java.util.concurrent.ScheduledFuture<?> outerFuture;

    public ScheduledFuture(Runnable runnable) {
        this.supplier = () -> {
            runnable.run();
            return null;
        };
    }

    public ScheduledFuture(Supplier<V> supplier) {
        this.supplier = supplier;
    }

    @Override
    public void run() {
        synchronized (lock) {
            if (state == 3) {
                // cancelled
                return;
            }
            assert state == 0;
            state = 1;
        }
        try {
            runningOn = Thread.currentThread();
            result = supplier.get();
        } catch (Throwable throwable) {
            this.throwable = throwable;
        } finally {
            runningOn = null;
            outerFuture = null;
            synchronized (lock) {
                assert state == 1;
                state = 2;
                lock.notifyAll();
            }
        }
    }

    @Override
    public boolean cancel(boolean mayInterruptIfRunning) {
        synchronized (lock) {
            if (state == 0) {
                state = 3;
                outerFuture.cancel(false);
                outerFuture = null;
                lock.notifyAll();
                return true;
            } else if (state == 1) {
                if (mayInterruptIfRunning) {
                    runningOn.interrupt();
                }
                return false;
            } else if (state == 2) {
                return false;
            } else if (state == 3) {
                return true;
            } else {
                assert false;
                return false;
            }
        }
    }

    @Override
    public boolean isCancelled() {
        return state == 3;
    }

    @Override
    public boolean isDone() {
        return state == 2;
    }

    @Override
    public V get() throws InterruptedException, ExecutionException {
        synchronized (lock) {
            while (true) {
                if (state == 2) {
                    if (throwable != null) {
                        throw new ExecutionException(throwable);
                    }
                    return result;
                } else if (state == 3) {
                    return null;
                } else if (state == 1) {
                    lock.wait();
                } else {
                    assert false;
                    return null;
                }
            }
        }
    }

    @Override
    public V get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
        synchronized (lock) {
            while (true) {
                if (state == 2) {
                    if (throwable != null) {
                        throw new ExecutionException(throwable);
                    }
                    return result;
                } else if (state == 3) {
                    return null;
                } else if (state == 1) {
                    lock.wait(unit.toMillis(timeout));
                    if (state == 1) {
                        throw new TimeoutException();
                    }
                } else {
                    assert false;
                    return null;
                }
            }
        }
    }
}
